from scipy.io import loadmat
from PIL import Image
import numpy as np
import os
from glob import glob
import cv2
import argparse

import xml.etree.ElementTree as ET


def get_points(label):
    tree = ET.parse(label)
    root = tree.getroot()
    points = [[int(obj.find('point').find('x').text), int(obj.find('point').find('y').text)]for obj in root.findall('object')]
    return points
    
    
def cal_new_size(im_h, im_w, min_size, max_size):
    if im_h < im_w:
        if im_h < min_size:
            ratio = 1.0 * min_size / im_h
            im_h = min_size
            im_w = round(im_w*ratio)
        elif im_h > max_size:
            ratio = 1.0 * max_size / im_h
            im_h = max_size
            im_w = round(im_w*ratio)
        else:
            ratio = 1.0
    else:
        if im_w < min_size:
            ratio = 1.0 * min_size / im_w
            im_w = min_size
            im_h = round(im_h*ratio)
        elif im_w > max_size:
            ratio = 1.0 * max_size / im_w
            im_w = max_size
            im_h = round(im_h*ratio)
        else:
            ratio = 1.0
    return im_h, im_w, ratio


def find_dis(point):
    square = np.sum(point*points, axis=1)
    dis = np.sqrt(np.maximum(square[:, None] - 2*np.matmul(point, point.T) + square[None, :], 0.0))
    dis = np.mean(np.partition(dis, 3, axis=1)[:, 1:4], axis=1, keepdims=True)
    return dis

def generate_data(im_path):
    
    im = Image.open(im_path)
    im_w, im_h = im.size
    
    print(f"im_path=={im_path}, im.size={im.size}")
    
    gt_path = im_path.replace('/Infrared/', '/GT_/').replace('R.jpg', 'R.xml')
    
    points = get_points(gt_path)
    points = np.asarray(points)
    
    idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
    points = points[idx_mask]
    im_h, im_w, rr = cal_new_size(im_h, im_w, min_size, max_size)
    im = np.array(im)
    if rr != 1.0:
        im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
        points = points * rr
    return Image.fromarray(im), points


def parse_args():
    parser = argparse.ArgumentParser(description='Test ')
    parser.add_argument('--origin-dir', default='/data/store1/nzd/align/test_prj/datasets/DroneRGBT/',
                        help='original data directory')
    parser.add_argument('--data-dir', default='./drone',
                        help='processed data directory')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    save_dir = args.data_dir
    min_size = 512
    max_size = 2048

    for phase in ['Train', 'Test']:
        sub_dir = os.path.join(args.origin_dir, phase)
        if phase == 'Train':
            sub_phase = 'train'
            sub_save_dir = os.path.join(save_dir, sub_phase)
            if not os.path.exists(sub_save_dir):
                os.makedirs(sub_save_dir)
            
            with open('{}.txt'.format(sub_phase)) as f:
                gt_list = glob(os.path.join(sub_dir, 'GT_/*.xml'))
                
                for gt in gt_list:
                    im_path = gt.replace('/GT_/', '/Infrared/').replace('R.xml', 'R.jpg')
                    name = os.path.basename(gt)

                    print(name)

                    if name in ['636R.xml', '656R.xml']:
                        continue

                    im, points = generate_data(im_path)
                    
                    if sub_phase == 'train':
                        dis = find_dis(points)
                        points = np.concatenate((points, dis), axis=1)
                    
                    im_save_path = os.path.join(sub_save_dir, name.replace('.xml', '.jpg'))
                    im.save(im_save_path)
                    gd_save_path = im_save_path.replace('jpg', 'npy')
                    np.save(gd_save_path, points)
        else:
            sub_save_dir = os.path.join(save_dir, 'test')
            if not os.path.exists(sub_save_dir):
                os.makedirs(sub_save_dir)
            gt_list = glob(os.path.join(sub_dir, 'GT_/*.xml'))
            for gt in gt_list:
                im_path = gt.replace('/GT_/', '/Infrared/').replace('R.xml', 'R.jpg')
                name = os.path.basename(gt)
                #print(name)
                im, points = generate_data(im_path)
                im_save_path = os.path.join(sub_save_dir, name.replace('.xml', '.jpg'))
                im.save(im_save_path)
                gd_save_path = im_save_path.replace('jpg', 'npy')
                np.save(gd_save_path, points)
